from manim import *

COLORS = (GREEN, RED, BLUE)

class MyScene(Scene):
    def small_pause(self, n=0.5):
        self.wait(n)

    def pause(self, n=1.5):
        self.wait(n)

    def medium_pause(self, n=3):
        self.wait(n)

    def long_pause(self, n=5):
        self.wait(n)


# 给矩阵设置颜色
def matrix_set_color(left_matrix, right_matrix):
    left_matrix.get_rows()[0].set_color(COLORS[0])
    left_matrix.get_rows()[1].set_color(COLORS[1])
    left_matrix.get_rows()[2].set_color(COLORS[2])
    right_matrix.get_rows()[0].set_color(COLORS[0])
    right_matrix.get_rows()[1].set_color(COLORS[1])
    right_matrix.get_rows()[2].set_color(COLORS[2])


class EliminationMethod(MyScene):
    def construct(self):

        # 原始线性方程组
        eq1 = VGroup(
            MathTex(*r"1  \square  +  1  \triangle  +  0  \lozenge  =  3".split('  ')),
            MathTex(*r"2  \square  +  0  \triangle  +  1  \lozenge  =  2".split('  ')),
            MathTex(*r"0  \square  +  1  \triangle  +  1  \lozenge  =  -1".split('  ')),
        ).arrange(DOWN)
        eq1.shift(UP*1.5 + LEFT*2)
        # 设置颜色
        for i in [0, 3, 6, -1]:
            eq1[0][i].set_color(COLORS[0])
            eq1[1][i].set_color(COLORS[1])
            eq1[2][i].set_color(COLORS[2])
        # 系数A
        coefficient = VGroup(
            eq1[0][0], eq1[0][3], eq1[0][6],
            eq1[1][0], eq1[1][3], eq1[1][6],
            eq1[2][0], eq1[2][3], eq1[2][6]
        )
        # 未知数X
        unknown = VGroup(
            eq1[0][1], eq1[0][4], eq1[0][7],
            eq1[1][1], eq1[1][4], eq1[1][7],
            eq1[2][1], eq1[2][4], eq1[2][7],
        )
        # 相等符号=
        equal = VGroup(
            eq1[0][-2], eq1[1][-2], eq1[2][-2]
        )
        # 右边b
        right = VGroup(
            eq1[0][-1], eq1[1][-1], eq1[2][-1]
        )

        # 用矩阵语言描述
        left_matrix = Matrix([[1, 1, 0], [2, 0, 1], [0, 1, 1]])
        unknown_matrix = Matrix([["\square"], ["\\triangle"], ["\lozenge"]])
        equal_matrix = MathTex("=")
        right_matrix = Matrix([[3], [2], [-1]])
        matrix = VGroup(left_matrix, unknown_matrix, equal_matrix, right_matrix).arrange(RIGHT)
        matrix.shift(DOWN*1.5 + LEFT*2)

        # 设置颜色
        matrix_set_color(left_matrix, right_matrix)

        # 问题
        question = MathTex(r"\begin{cases} \square &= ? \\ \triangle &= ? \\ \lozenge &= ? \end{cases}")
        question.shift(UP + RIGHT*3)

        # Scene2
        # 增广矩阵
        augmented_matrix = VGroup(left_matrix.copy(), right_matrix.copy()).arrange(RIGHT)
        augmented_matrix.shift(UP*1.5 + LEFT*4)

        # 第2行第1列:中间计算过程
        temp_21 = VGroup(
            MathTex(*r"2  -  1  \times  2  =  0".split('  ')),
            MathTex(*r"0  -  1  \times  2  =  -2".split('  ')),
            MathTex(*r"1  -  0  \times  2  =  1".split('  ')),
            MathTex(*r"2  -  3  \times  2  =  -4".split('  ')),
            Arrow(start=LEFT*2, end=RIGHT*2, stroke_width=2, max_tip_length_to_length_ratio=0.1),
            Text("第2行第1列", color=YELLOW)
        ).arrange(DOWN, aligned_edge=LEFT).scale(0.5)
        temp_21.next_to(augmented_matrix, RIGHT*2.5)
        # 给中间计算过程设置颜色
        for i in range(0, 4, 1):
            temp_21[i][0].set_color(COLORS[1])
            temp_21[i][2].set_color(COLORS[0])
            temp_21[i][-1].set_color(COLORS[1])

        # 第2行第1列: 结果
        left_21 = Matrix([[1, 1, 0], [0, -2, 1], [0, 1, 1]])
        right_21 = Matrix([[3], [-4], [-1]])
        left_21.shift(UP*1.5 + RIGHT*3)
        right_21.next_to(left_21, RIGHT)
        right_21.add_updater(lambda x : x.next_to(left_21, RIGHT))
        # 设置颜色
        matrix_set_color(left_21, right_21)


        # 第3行第2列: 中间计算过程
        temp_32 = VGroup(
            MathTex(*r"0  \times  2  +  0  =  0".split('  ')),
            MathTex(*r"1  \times  2  +  (-2)  =  0".split('  ')),
            MathTex(*r"1  \times  2  +  1  =  3".split('  ')),
            MathTex(*r"(-1)  \times  2  +  (-4)  =  -6".split('  ')),
            Arrow(start=LEFT*2, end=RIGHT*2.5, stroke_width=2, max_tip_length_to_length_ratio=0.1),
            Text("第3行第2列", color=YELLOW)
        ).arrange(DOWN, aligned_edge=LEFT).scale(0.5)
        temp_32.next_to(temp_21, DOWN*4).shift(LEFT*0.4)
        # 给中间计算过程设置颜色
        for i in range(0, 4, 1):
            temp_32[i][0].set_color(COLORS[2])
            temp_32[i][4].set_color(COLORS[1])
            temp_32[i][-1].set_color(COLORS[2])
        # 第3行第2列: 结果
        left_32 = Matrix([[1, 1, 0], [0, -2, 1], [0, 0, 3]])
        right_32 = Matrix([[3], [-4], [-6]])
        left_32.shift(DOWN*1.5 + RIGHT*3)
        right_32.next_to(left_32, RIGHT)
        right_32.add_updater(lambda x : x.next_to(left_32, RIGHT))
        # 设置颜色
        matrix_set_color(left_32, right_32)

        # 消元目标背景框 
        target1_box = SurroundingRectangle(augmented_matrix[0][0][3])
        target2_box = SurroundingRectangle(left_21[0][7])


        # Scene3
        # 化简后的线性方程组
        eq2 = MathTex(*r"1  \square  +  1  \triangle  +  0  \lozenge  &=  3  \\  -2  \triangle  +  1  \lozenge  &=  -4  \\  3  \lozenge  &=  6".split('  '))
        eq2.shift(LEFT*4)
        for i in [0, 3, 6, 9]:
            eq2[i].set_color(COLORS[0])
        for i in [11, 14, 17]:
            eq2[i].set_color(COLORS[1])
        for i in [19, -1]:
            eq2[i].set_color(COLORS[2])

        # 箭头
        arrow1 = Arrow(start=DOWN, end=UP, path_arc=PI, stroke_width=1, max_tip_length_to_length_ratio=0.1)
        arrow1.scale(0.35)
        arrow1.move_to(eq2[22], RIGHT).shift(UP*0.3 + RIGHT*0.9)
        arrow2 = arrow1.copy()
        arrow2.move_to(eq2[17], RIGHT).shift(UP*0.3 + RIGHT*0.5)

        # 中间答案
        answer1 = MathTex(r"\lozenge = -2").next_to(arrow1, RIGHT)
        answer2 = MathTex(r"\triangle = 1").next_to(arrow2, RIGHT)
        answer3 = MathTex(r"\square = 2").next_to(answer2, UP)

        # 最后答案
        answer = MathTex(r"\begin{cases} \square &= 2 \\ \triangle &=  1 \\ \lozenge &= -2 \end{cases}")
        answer.shift(RIGHT*3)




        # Animation
        # Scene1
        # 线性方程组
        self.wait(3)
        self.play(Write(eq1[0]), run_time=2)
        self.wait()
        self.play(Write(eq1[1]), run_time=2)
        self.wait()
        self.play(Write(eq1[2]), run_time=2)
        self.wait()

        # 问题
        self.play(FadeIn(question))
        self.wait()
        self.play(
            Wiggle(question[0][7]),
            Wiggle(question[0][10]),
            Wiggle(question[0][13]),
            scale_value=1.5,
            rotation_angle=0.1 * TAU,
            run_time=2
        )
        self.wait()
        # 矩阵语言
        self.play(
            ReplacementTransform(coefficient.copy(), left_matrix[0]),
            ReplacementTransform(unknown.copy(), unknown_matrix[0]),
            ReplacementTransform(equal.copy(), equal_matrix),
            ReplacementTransform(right.copy(), right_matrix[0])
        )
        self.play(
            FadeIn(left_matrix[1:3]),
            FadeIn(unknown_matrix[1:3]),
            FadeIn(right_matrix[1:3])
        )
        self.wait()

        # Scene1--->Scene2
        self.play(FadeOut(eq1, question, unknown_matrix, equal_matrix, right_matrix))
        self.play(ReplacementTransform(left_matrix, augmented_matrix[0]))

        # Scene2: 消元
        # 消元: 第2行第1列
        self.play(
            Create(target1_box),
            Create(temp_21[-2]), Write(temp_21[-1]),
        )
        self.play(
            FadeIn(left_21[1:3]), 
            ReplacementTransform(left_matrix[0][0:3].copy(), left_21[0][0:3])
        )
        self.play(ReplacementTransform(left_matrix[0][0:6].copy(),temp_21[0:3]))
        self.small_pause()
        self.play(ReplacementTransform(VGroup(temp_21[0][-1], temp_21[1][-1], temp_21[2][-1]).copy(), left_21[0][3:6]))
        self.play(ReplacementTransform(left_matrix[0][6:9].copy(), left_21[0][6:9]))


        # 消元: 第3行第2列
        self.play(
            Create(target2_box),
            Create(temp_32[-2]), Write(temp_32[-1]),
        )
        self.play(
            FadeIn(left_32[1:3]),
            ReplacementTransform(left_21[0][0:6].copy(), left_32[0][0:6])
        )
        self.play(ReplacementTransform(left_21[0][3:9].copy(), temp_32[0:3]))
        self.small_pause()
        self.play(ReplacementTransform(VGroup(temp_32[0][-1], temp_32[1][-1], temp_32[2][-1]).copy(), left_32[0][6:9]))
        self.small_pause()

        # self.add(left_21, left_32)
        # 增广矩阵
        self.play(FadeIn(augmented_matrix[1]))
        self.play(
            ReplacementTransform(augmented_matrix[1][0][0:2].copy(), temp_21[3]),
            FadeIn(right_21[1:3])
        )
        self.small_pause()
        self.play(ReplacementTransform(temp_21[3].copy(), right_21[0][0:2]))
        self.small_pause()
        self.play(ReplacementTransform(augmented_matrix[1][0][2].copy(), right_21[0][2]))
        
        self.play(
            ReplacementTransform(right_21[0][0].copy(), right_32[0][0]),
            FadeIn(right_32[1:3]),
        )
        self.play(ReplacementTransform(right_21[0][1:3].copy(), temp_32[3]))
        self.small_pause()
        self.play(ReplacementTransform(temp_32[3].copy(), right_32[0][1:3]))

        self.small_pause()

        # Scene2 --> Scene3
        self.play(FadeOut(augmented_matrix, temp_21, left_21, right_21, temp_32, target1_box, target2_box))
        self.play(
            ReplacementTransform(left_32.get_rows()[0], eq2[0:8]),
            ReplacementTransform(left_32.get_rows()[1], eq2[11:16]),
            ReplacementTransform(left_32.get_rows()[2], eq2[19:21]),
            FadeIn(eq2[8], eq2[16], eq2[21]),
            FadeTransform(right_32.get_rows()[0], eq2[9]),
            FadeTransform(right_32.get_rows()[1], eq2[17]),
            FadeTransform(right_32.get_rows()[2], eq2[22]),
            FadeOut(left_32[1:3], right_32[1:3]),
            run_time=2
        )
        self.small_pause()

        # Scene3: 回代
        self.play(FadeIn(arrow1), Write(answer1))
        self.play(FadeIn(arrow2), Write(answer2))
        self.play(Write(answer3))

        self.play(FadeIn(answer[0][0:5]))
        self.play(FadeTransform(VGroup(answer1, answer2, answer3).copy(), answer[0][5:15]))
        self.wait(1)

